import torch
import torch.nn as nn
import torch.nn.functional as F


class CausalConv3d(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size=3, dilation=1, **kwargs):
        super().__init__()
        self.padding = kernel_size // 2
        self.frame_padding = (kernel_size - 1) * dilation
        self.net = nn.Conv3d(in_channels=in_channels,
                             out_channels=out_channels, kernel_size=kernel_size,
                             stride=1, padding=0, dilation=dilation, **kwargs)

    def forward(self, x):
        # x -> (b, c, t, h, w)
        padding = tuple([self.padding] * 4) + (self.frame_padding, 0)
        x = F.pad(x, padding)
        x = self.net(x)

        return x
